import numpy as np 
from matplotlib import pyplot as plt
from tools.restore_model_tools import restore_model, plot_naca001065_flow_field, DL_Euler_Equation_2D

# 需要对哪个level的网格进行加密？
level = 2

model = restore_model(model_path='/home/aistudio/work/NACA0010-65/level' + str(level) + '/2_PINN_training/Model/model-75000.pdparams')
plot_naca001065_flow_field(model)

'''------------------------------------------- Residual-based adaptive refinement ---------------------------------------------'''
import triangle as tr
from tools.mesh_refinement_tools import get_top_k_indices, plot_naca001065_add_points, insert_points_between_points, insert_midpoints

# level1
# airfoil_points 是原翼型坐标点，形状为 (205, 2)
airfoil_points = np.load("../0_Mesh_files/ordered_airfoil_points.npy")
farfield_points = np.load(r'../0_Mesh_files/farfield_points.npy')
internal_points = np.load(r'../0_Mesh_files/internal_points.npy')

# refined_airfoil_points 是细化后的翼型坐标点，形状为 (410, 2)
refined_airfoil_points = insert_points_between_points(airfoil_points)

# 然后得到加密前网格的所有边的信息，使用-e选项输出Edges的信息：
polygon = np.concatenate((airfoil_points, farfield_points, internal_points))

# 设置约束边界的索引，翼型边界和矩形边界
inner_airfoil_index = np.hstack((np.arange(0, airfoil_points.shape[0], 1).reshape(airfoil_points.shape[0], 1), 
                                  np.arange(1, airfoil_points.shape[0]+1, 1).reshape(airfoil_points.shape[0], 1)))
inner_airfoil_index[airfoil_points.shape[0]-1, 1]=0
outer_farfield_index = np.hstack((np.arange(0, farfield_points.shape[0], 1).reshape(farfield_points.shape[0], 1), 
                                  np.arange(1, farfield_points.shape[0] +1, 1).reshape(farfield_points.shape[0], 1)))
outer_farfield_index[farfield_points.shape[0]-1, 1]=0
outer_farfield_index = outer_farfield_index + inner_airfoil_index.shape[0]
inner_outer_index = np.vstack((inner_airfoil_index, outer_farfield_index))

# Delauney 三角化
A = dict(vertices = polygon, segments = inner_outer_index, holes=[[0.5, 0]])
B = tr.triangulate(A,'pe')
tr.comparev(plt, A, B, figsize=(60, 45))
plt.show()

# 节点和边的信息
edges = B['edges']
vertices = B['vertices']

# 筛选出原网格除了翼型的边之外的所有边：
connected_edges = set()
for edge_index in edges:
    vertex1_index, vertex2_index = edge_index
    if vertex1_index in inner_airfoil_index and vertex2_index in inner_airfoil_index:
        connected_edges.add(tuple(edge_index))  # Convert the numpy.ndarray to a tuple
filtered_edges = [edge_index for edge_index in edges if tuple(edge_index) not in connected_edges]

# 得到除翼型外所有Edges的中点：
inner_midpoints = []
for edge in filtered_edges:
    vertex1 = vertices[edge[0]]
    vertex2 = vertices[edge[1]]
    inner_midpoint = (vertex1 + vertex2) / 2
    inner_midpoints.append(inner_midpoint)
inner_midpoints = np.array(inner_midpoints)

# 合并所有边的中点，翼型点在前，其他点在后：
all_edges_midpoints = np.vstack((refined_airfoil_points[1::2], inner_midpoints))

# 最后我们使用原网格的Edges中点的13168个区域坐标点去计算PDEs的残差，并从中筛选出前2000个左右（约前15%）对应PDEs残差最大的点并绘制成散点图
X = all_edges_midpoints
[f1, f2, f3, f4] = model.predict(X, operator=DL_Euler_Equation_2D)
err_eq = np.absolute(f1) + np.absolute(f2) + np.absolute(f3) + np.absolute(f4)
x_id = get_top_k_indices(err_eq, k=1896)
add_points = X[x_id]
plot_naca001065_add_points(X[x_id])

# 更新新网格的所有点
# 处于翼型中点的那些点
airfoil_midpoints = X[x_id[x_id < refined_airfoil_points[1::2].shape[0]]]
# 翼型上待加密的点插到原翼型所有点里
airfoil_newpoints = insert_midpoints(airfoil_midpoints, airfoil_points)
# 不包括原网格的节点的中点
internal_midpoints =  X[x_id[x_id >= refined_airfoil_points[1::2].shape[0]]]

'''----------------------------- Delaunay triangulation for refinement mesh--------------------------'''
polygon = np.concatenate((airfoil_newpoints, farfield_points, internal_points, internal_midpoints))

# 设置约束边界的索引，翼型边界和矩形边界
inner_airfoil_index = np.hstack((np.arange(0, airfoil_newpoints.shape[0], 1).reshape(airfoil_newpoints.shape[0], 1), 
                                  np.arange(1, airfoil_newpoints.shape[0]+1, 1).reshape(airfoil_newpoints.shape[0], 1)))
inner_airfoil_index[airfoil_newpoints.shape[0]-1, 1]=0
outer_farfield_index = np.hstack((np.arange(0, farfield_points.shape[0], 1).reshape(farfield_points.shape[0], 1), 
                                  np.arange(1, farfield_points.shape[0] +1, 1).reshape(farfield_points.shape[0], 1)))
outer_farfield_index[farfield_points.shape[0]-1, 1]=0
outer_farfield_index = outer_farfield_index + inner_airfoil_index.shape[0]
inner_outer_index = np.vstack((inner_airfoil_index, outer_farfield_index))

# # 将边界连接信息添加到字典 A 中

C = dict(vertices = polygon, segments = inner_outer_index, holes=[[0.5, 0]])
D = tr.triangulate(C,'p')
tr.comparev(plt, C, D, figsize=(60, 45))
plt.show()

'''----------------------------- convert midpoint_refinement mesh to vtk and msh --------------------------'''
import meshio

# 获取三角剖分结果的顶点和单元格信息
vertices = D['vertices']
cells = [("triangle", D['triangles'])]

# 创建meshio的网格对象
mesh = meshio.Mesh(points=vertices, cells=cells)

# 定义保存的文件名
vtkFileName = '../../level' + str(level+1) + '/0_Mesh_files/level' + str(level+1) + '_midpoints_refinement_mesh.vtk'
# 将网格对象保存为VTK文件
meshio.write(vtkFileName, mesh, file_format="vtk", binary=False)

# 使用meshio读取VTK文件
mesh = meshio.read(vtkFileName)

# 将网格对象保存为Gmsh支持的格式（例如.msh文件）
mshFileName = '../../level' + str(level+1) + '/0_Mesh_files/level' + str(level+1) + '_midpoints_refinement_mesh.msh'
meshio.write(mshFileName, mesh, file_format='gmsh22',binary=False)
print("转换后的文件已保存为:", mshFileName)


'''----------------------------- refinement mesh nodes --------------------------'''
np.save(r'../../level' + str(level+1) + '/0_Mesh_files/ordered_airfoil_points.npy', airfoil_newpoints)
np.save(r'../../level' + str(level+1) + '/0_Mesh_files/farfield_points.npy', farfield_points)
np.save(r'../../level' + str(level+1) + '/0_Mesh_files/internal_points.npy', np.concatenate((internal_points, internal_midpoints)))

